"""
==========================================
量子傅里叶变换 (Quantum Fourier Transform, QFT)
==========================================

本文件实现了 QFT 与其逆变换 IQFT（Inverse QFT），适用于 cqlib 的 Circuit 接口。

统一入口：
    fourier_on(circuit, qubits, inverse=False, do_swaps=True) -> Circuit
    build_fourier(n_qubits, inverse=False, do_swaps=True) -> Circuit

主要思想：
---------
量子傅里叶变换是经典傅里叶变换的量子化版本，定义为：

    QFT|x> = (1/√N) Σ_y exp(2πi * x*y / N) |y>

它在量子算法（如 Shor 分解算法、相位估计算法）中扮演核心角色。

实现思路：
---------
- 每个量子比特需要经历一个 Hadamard 门和一系列受控相移门（Controlled Phase, CP）。
- 对第 i 个量子比特，受控角度为 π / 2^(j - i)，其中 j > i。
- 若 `inverse=True`，则相位角符号取反，门序颠倒。
- 可选的 swap 操作用于纠正位序（QFT 输出的顺序与输入相反）。

物理意义：
----------
- QFT 把计算基态 |x> 变换为频域叠加态；
- IQFT 则用于把频域信息还原回原始基态；
- QFT 不依赖测量，属于完全幺正变换。

注意：
------
- 本实现中 CP(θ) 门使用 RZ 与 CRZ 的等效分解实现：
      CP(θ) ≈ RZ(control, θ/2) · CRZ(control->target, θ)
  差一个全局相位，但不影响物理结果与测试比对。
- cqlib 的 RZ 接口为 rz(qubit, theta)（单参数形式）。

"""

from math import pi
from typing import Iterable, List

from circuits.algorithm_circuit import Algorithm_Circuit as Circuit


# ============================================================
# 辅助函数：受控相移的分解 (CP → RZ + CRZ)
# ============================================================
def _cp_via_crz(qc: Circuit, control: int, target: int, theta: float) -> None:
    """
    构造受控相移门 CP(θ)。

    数学形式：
        CP(θ) = diag(1, 1, 1, e^{iθ})
    实现等价：
        CP(θ) = RZ(control, θ/2) · CRZ(control->target, θ)
    （仅忽略全局相位）

    参数：
        qc:      电路对象
        control: 控制量子比特索引
        target:  目标量子比特索引
        theta:   相移角度（弧度）
    """
    qc.rz(control, theta / 2)
    qc.crz(control, target, theta)


# ============================================================
# 主函数：量子傅里叶变换 QFT / IQFT
# ============================================================
def qft(circuit: Circuit, qubits: Iterable[int],
        inverse: bool = False, do_swaps: bool = True) -> Circuit:
    """
    在给定电路上对指定量子比特集合执行 QFT 或其逆变换。

    参数：
        circuit : Circuit
            已存在的量子电路对象
        qubits  : Iterable[int]
            参与 QFT 的量子比特索引序列（从低位到高位）
        inverse : bool
            若为 True，执行逆 QFT（IQFT）；否则执行正向 QFT
        do_swaps : bool
            是否执行位反转 (bit-reversal) 交换操作

    返回：
        circuit : Circuit
            修改后的电路对象

    -------------------------
    实现细节：
    -------------------------
    - QFT 顺序（正向）：
        1. 对每个量子位 i：
            a. 施加 Hadamard 门；
            b. 对所有更高位 j>i，施加受控相移 CP(π / 2^(j-i))
        2. 若 do_swaps=True，则执行比特序逆转（swap i ↔ n-1-i）

    - IQFT 顺序（逆向）：
        1. 若 do_swaps=True，先执行比特逆转；
        2. 对每个量子位 i（从高位到低位）：
            a. 对所有更高位 j>i，施加受控相移 CP(-π / 2^(j-i))
            b. 再施加 Hadamard 门。

    注：
        - “高位控制、低位目标” 的结构可保持一致。
        - θ 取负号即可实现逆傅里叶变换。
    """
    q: List[int] = list(qubits)
    n = len(q)

    if not inverse:
        # ====== 正向 QFT ======
        for i in range(n):
            # Step 1: Hadamard
            circuit.h(q[i])

            # Step 2: Controlled-phase (CP) rotations
            for j in range(i + 1, n):
                theta = pi / (2 ** (j - i))
                _cp_via_crz(circuit, q[j], q[i], theta)

        # Step 3: 可选的比特序反转（纠正位序）
        if do_swaps:
            for i in range(n // 2):
                circuit.swap(q[i], q[n - 1 - i])

    else:
        # ====== 逆向 QFT ======
        # Step 1: 若需要，先反转位序
        if do_swaps:
            for i in range(n // 2):
                circuit.swap(q[i], q[n - 1 - i])

        # Step 2: 从高位到低位逆序施加门操作
        for i in reversed(range(n)):
            # 先施加逆向的受控相移（符号取负）
            for j in reversed(range(i + 1, n)):
                theta = -pi / (2 ** (j - i))
                _cp_via_crz(circuit, q[j], q[i], theta)
            # 最后施加 Hadamard
            circuit.h(q[i])

    return circuit


# ============================================================
# 构建独立 QFT 电路的便捷函数
# ============================================================
def build_qft(n_qubits: int, inverse: bool = False, do_swaps: bool = True) -> Circuit:
    """
    构建独立的 QFT 或 IQFT 电路（不依赖外部 Circuit）。

    参数：
        n_qubits : int
            量子比特数量
        inverse  : bool
            是否构建逆变换 IQFT
        do_swaps : bool
            是否执行比特反转操作

    返回：
        qc : Circuit
            完整的 QFT/IQFT 电路对象

    示例：
        >>> qc = build_qft(3)
        >>> qc.draw()
        （输出三量子比特 QFT 电路图）
    """
    qc = Circuit(n_qubits)
    return qft(qc, range(n_qubits), inverse=inverse, do_swaps=do_swaps)
